#!/usr/bin/python

__author__ = "Jimmy Saw"
__copyright__ = "Copyright 2011, My Auto-annotator"
__credits__ = "Jimmy Saw"
__email__ = "jimmysaw@gmail.com"

"""
TO DO: 
Usage: python auto_anno.py annofile.txt seqfile.fasta

"""

import sys
import re
from Bio import SeqIO

start_codons = {'ATG': 'Start', 'CTG': 'Start', 'GTG': 'Start', 'TTG': 'Start'}
stop_codons = {'TAA': 'Stop', 'TAG': 'Stop', 'TGA': 'Stop'}

#Formatting for 4 significant figures (for decimals)
def fmt(f):
    st = '{0:.4}'.format(f)
    return st

#Indicator
def indicate():
    st = "<-------------------------------------"
    return st

#Check start sites begins
annofile = sys.argv[1]

seqfile = sys.argv[2]

af = open(annofile, "rU")

sf = open(seqfile, "rU")

rec = SeqIO.read(sf, "fasta")

lines = af.readlines()

num = len(lines)

i = 0

good_count = 0

bad_count = 0

has_blast = 0

tofix_list = {}

while i < num:
    curr_line = lines[i].split('\t')
    curr_prodigal_id = curr_line[0]
    curr_locus_tag = curr_line[1]
    curr_feat_type = curr_line[2]
    curr_start = int(curr_line[3])
    curr_stop = int(curr_line[4])
    curr_frame = curr_line[5]
    curr_blast_status = curr_line[6]
    curr_blast_evalue = curr_line[7]
    curr_tophit_id = curr_line[8]
    curr_tophit_sim = curr_line[9]
    curr_tophit_acc = curr_line[10]
    curr_tophit_gi = curr_line[11]
    curr_blast_qstart = curr_line[12]
    curr_blast_qstop = curr_line[13]
    curr_hitlength = curr_line[14]
    curr_blast_mstart = curr_line[15]
    curr_blast_mstop = curr_line[16]
    curr_hit_def = curr_line[17]
    curr_hit_org = curr_line[18]
    curr_COG_status = curr_line[19]
    curr_COG_len = curr_line[20]
    curr_COG_evalue = curr_line[21]
    curr_COG_pid = curr_line[22]
    curr_COG_psim = curr_line[23]
    curr_COG_start = curr_line[24]
    curr_COG_stop = curr_line[25]
    curr_COG_id = curr_line[26]
    curr_COG_cat = curr_line[27]
    curr_COG_def = curr_line[28]
    curr_gene_name = curr_line[29]
    curr_gene_name_fixed = curr_line[30]
    curr_cluster_def = curr_line[31]
    curr_cluster_org = curr_line[32]

    if curr_feat_type == "rRNA":
    #print curr_locus_tag, "\t", curr_hit_def
        print curr_locus_tag, "\t", curr_feat_type, "\t", curr_hit_def
    elif curr_feat_type == "tRNA":
    #print curr_locus_tag, "\t", curr_hit_def
        print curr_locus_tag, "\t", curr_feat_type, "\t", curr_hit_def
    elif curr_feat_type == "CDS":
    #Work on the coding sequences for start site correction
        if curr_blast_status == "Yes" and int(curr_tophit_id) >= 75:
            has_blast += 1
            curr_aa_len = (curr_stop - curr_start + 1 - 3)/3 #minus 3 for stop codon
            curr_coverage = (curr_aa_len / float(curr_hitlength))*100
            blastfile = curr_prodigal_id + ".blastp.tbl"
            f = open(blastfile, "rU")
            bfl = f.readlines()
            ln = len(bfl)
            if ln == 1:
                #Check length
                mstart = int(curr_blast_mstart)
                qstart = curr_start
                if qstart == mstart:
                    good_count += 1
                    print curr_prodigal_id, "\t", "No need to correct start site"
                else:
                    bad_count += 1
                    if qstart < mstart:
                        fixlen = (mstart - qstart) * 3
                        fixtype = "extend"
                        tofix_list[curr_prodigal_id] = "Need to extend start site"
                        print curr_prodigal_id, "\t", "Need to extend start site", "\t", fixlen
                    else:
                        fixlen = (qstart - mstart) * 3
                        fixtype = "cut"
                        tofix_list[curr_prodigal_id] = "Need to shorten start site"
                        print curr_prodigal_id, "\t", "Need to shorten start site", "\t", fixlen

            elif 1 < ln < 10: # Blast hits less than 10 but more than 1
                vote = 0
                length_vote = 0
                coverage_vote = 0
                k = 0
                to_cut = []
                to_extend = []
                while k < ln:
                    cl = bfl[k].split('\t')
                    qlen = float(cl[6])
                    qstart = int(cl[7])
                    qstop = int(cl[8])
                    mlen = float(cl[10])
                    mstart = int(cl[11])
                    mstop = int(cl[12])
                    coverage = (qlen/mlen) * 100
                    if qstart == mstart:
                        vote += 1
                    else:
                        if qstart < mstart: # too short
                            to_extend.append(mstart)
                        else: # too long
                            to_cut.append(qstart)
                    k += 1

                if vote > ln/2:
                    good_count += 1
                    print curr_prodigal_id, "\t", "No need to correct start site!" 
                else:
                    tofix_list[curr_prodigal_id] = "Need to fix (n = " + str(ln) + ")"
                    bad_count += 1
                    if len(to_extend) == len(to_cut): # if equal votes, don't fix it
                        print curr_prodigal_id, "\t", "Not enough evidence to cut or extend"
                    elif len(to_extend) > len(to_cut):
                        uniq_to_extend = set(to_extend)
                        uniq_counts = {}
                        for x in uniq_to_extend:
                            uniq_counts[x] = to_extend.count(x)
                        largest = max(uniq_counts, key = lambda a: uniq_counts.get(a))
                        smallest = min(uniq_counts, key = lambda a: uniq_counts.get(a))
                        if largest > 1:
                            fixlen = (largest - int(curr_blast_qstart)) * 3
                            print curr_prodigal_id, "\t", "Need to extend start site"
                        else:
                            fixlen = 0
                            print curr_prodigal_id, "\t", "Need to extend start site but shouldn't. So, left it alone."
                    else:
                        uniq_to_cut = set(to_cut)
                        uniq_counts = {}
                        for x in uniq_to_cut:
                            uniq_counts[x] = to_cut.count(x)
                        largest = max(uniq_counts, key = lambda a: uniq_counts.get(a))
                        smallest = min(uniq_counts, key = lambda a: uniq_counts.get(a))
                        if largest > 1:
                            fixlen = (largest - int(curr_blast_qstart)) * 3
                            print curr_prodigal_id, "\t", "Need to shorten start site", "\t", fixlen
                        else:
                            fixlen = 0
                            print curr_prodigal_id, "\t", "Need to shorten start site but shouldn't. So, left it alone."

            else: # Blast hits >= 10
                vote = 0
                k = 0
                to_cut = []
                to_extend = []
                while k < 10:
                    cl = bfl[k].split('\t')
                    mstart = int(cl[11])
                    qstart = int(cl[7])
                    if qstart == mstart:
                        vote += 1
                    else:
                        if qstart < mstart: # too short
                            to_extend.append(mstart)
                        else: # too long
                            to_cut.append(qstart)
                    k += 1
                if vote >= 5:
                    good_count += 1
                    print curr_prodigal_id, "\t", vote, "\t", "No need to correct start site!"
                else:
                    tofix_list[curr_prodigal_id] = "Need to fix (n >= 10)"
                    bad_count += 1
                    if len(to_cut) == len(to_extend): # if equal votes, don't fix it
                        print curr_prodigal_id, "\t", "Not enough evidence to correct start site"
                    elif len(to_extend) > len(to_cut):
                        uniq_to_extend = set(to_extend)
                        uniq_counts = {}
                        for x in uniq_to_extend:
                            uniq_counts[x] = to_extend.count(x)
                        largest = max(uniq_counts, key = lambda a: uniq_counts.get(a))
                        smallest = min(uniq_counts, key = lambda a: uniq_counts.get(a))
                        if largest > 1:
                            fixlen = (largest - int(curr_blast_qstart)) * 3
                            print curr_prodigal_id, "\t", "Need to extend start site", "\t", fixlen 
                        else:
                            fixlen = 0
                            print curr_prodigal_id, "\t", "Need to extend but shouldn't. So, left it alone."
                    else:
                        uniq_to_cut = set(to_cut)
                        uniq_counts = {}
                        for x in uniq_to_cut:
                            uniq_counts[x] = to_cut.count(x)
                        largest = max(uniq_counts, key = lambda a: uniq_counts.get(a))
                        smallest = min(uniq_counts, key = lambda a: uniq_counts.get(a))
                        if largest > 1:
                            fixlen = (largest - int(curr_blast_qstart)) * 3
                            print curr_prodigal_id, "\t", "Need to shorten start site" , "\t", fixlen
                        else:
                            fixlen = 0
                            print curr_prodigal_id, "\t", "Need to cut but shouldn't. So, left it alone."
                    #print curr_prodigal_id, "\t", vote, "\t", "Need to fix!", fixstart(curr_prodigal_id)
        elif curr_blast_status == "Yes" and int(curr_tophit_id) < 75:
            has_blast += 1
            good_count += 1
            print curr_prodigal_id, "\t", "No need to fix. Top Hit AA Identity < 65%"
         
        else:
        #assign definition as "hypothetical protein"
        #done checking for start-site correction
            curr_nt_len = curr_stop - curr_start + 1
            if curr_nt_len < 90:
                print curr_prodigal_id, "\t", "Needs to be removed!", indicate()
            else:
                print curr_prodigal_id, "\t", "Hypothetical protein. Not fixed"

    else:
        print "Unknown feature type!"
    i += 1

total_count = good_count + bad_count

print "Total good counts:", "\t", good_count
print "Total bad counts:", "\t", bad_count
print "Total count:", "\t", total_count
print "Total count with BLAST hits:", "\t", has_blast
print "Total bad records to fix:", "\t", len(tofix_list)

    #Correct start sites

#print rec.id, len(rec), rec.seq[0:103]

for k, v in tofix_list.items():
    print k, "\t", v

    #Assignment of annotations

af.close()
sf.close()
